#pragma GCC optimize(2)
#pragma GCC optimize(3,"Ofast","inline")

#include<iostream>
#include "arg_max_with_value_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
const uint32_t BLOCK_SIZE = 32;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    ArgMaxWithValueTilingData tiling;
    auto dt = context->GetInputTensor(0)->GetDataType();
    uint32_t datatype;
    if(dt == ge::DT_UINT8){  // 8 bit
        datatype = 1;  // UINT8
    }else if(dt == ge::DT_FLOAT16 || dt == ge::DT_BF16){ // 16 bit
        datatype = 2; // float16
    }else if(dt == ge::DT_INT32){
        datatype = 3; // int32
    }
    else{
        datatype = 4; // float32
    }
    auto dimension = *context->GetAttrs()->GetInt(0);
    tiling.set_data_type(datatype);
    const gert::Shape vec = context->GetInputShape(0)->GetOriginShape();
    uint32_t shapeNum = vec.GetDimNum();
    if(datatype == 2){
        context->SetBlockDim(1);
        tiling.set_block_size(vec.GetDim(0));
        tiling.set_nxt_jump(vec.GetDim(1));
        tiling.set_nxt_idx(1);
    }else if(datatype == 3){
        context->SetBlockDim(1);
        tiling.set_block_size(vec.GetDim(0) * vec.GetDim(1));
        tiling.set_nxt_jump(1);
        tiling.set_nxt_idx(vec.GetDim(2));
    }else if(datatype == 1){
        context->SetBlockDim(1);
        tiling.set_block_size(vec.GetDim(0) * vec.GetDim(1) * vec.GetDim(2));
        tiling.set_nxt_jump(vec.GetDim(1) * vec.GetDim(2));
        tiling.set_nxt_idx(1);
    }else{
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        uint32_t aivNum = ascendcPlatform.GetCoreNum();
        uint32_t M = vec.GetDim(0) , N = vec.GetDim(1);
        uint32_t align = BLOCK_SIZE / 4;
        uint32_t block_size = (M + align - 1) / align; 
        uint32_t core = (block_size + aivNum - 1) / aivNum * align; 
        uint32_t aiv = (M + core - 1) / core; 
        
        context->SetBlockDim(aiv);
        // std::cout<< "aivNum : " << aivNum << std::endl; 
        tiling.set_block_size(core);
        tiling.set_nxt_jump(N);
        tiling.set_nxt_idx(M);
    }

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class ArgMaxWithValue : public OpDef {
public:
    explicit ArgMaxWithValue(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_UINT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("indice")
            .ParamType(REQUIRED)
            .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("values")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32, ge::DT_UINT8})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("dimension").Int();
        this->Attr("keep_dims").AttrType(OPTIONAL).Bool(false);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(ArgMaxWithValue);
}
